# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Quaternion
"""

from typing import Union, Tuple
import numpy as np
from numpy import ndarray
import torch
from torch import Tensor

__all__ = [
    'quaternion_to_rotation_matrix',
    'rotation_matrix_to_quaternion',
    'quaternion_multiply',
    'quaternion_conjugate',
    'quaternion_inverse',
    'quaternion_normalize',
    'quaternion_rotate',
]

def quaternion_to_rotation_matrix(quaternion: Union[Tensor, ndarray]) -> Union[Tensor, ndarray]:
    r"""
    Convert quaternion to rotation matrix.

    Args:
        quaternion (Union[Tensor, ndarray]): Quaternion with shape (..., 4).

    Returns:
        Union[Tensor, ndarray], rotation matrix with shape (..., 3, 3).
    """
    if isinstance(quaternion, ndarray):
        quaternion = torch.from_numpy(quaternion)
    
    q0, q1, q2, q3 = torch.unbind(quaternion, dim=-1)
    
    r00 = 1 - 2 * (q2 * q2 + q3 * q3)
    r01 = 2 * (q1 * q2 - q0 * q3)
    r02 = 2 * (q1 * q3 + q0 * q2)
    
    r10 = 2 * (q1 * q2 + q0 * q3)
    r11 = 1 - 2 * (q1 * q1 + q3 * q3)
    r12 = 2 * (q2 * q3 - q0 * q1)
    
    r20 = 2 * (q1 * q3 - q0 * q2)
    r21 = 2 * (q2 * q3 + q0 * q1)
    r22 = 1 - 2 * (q1 * q1 + q2 * q2)
    
    rotation_matrix = torch.stack([
        torch.stack([r00, r01, r02], dim=-1),
        torch.stack([r10, r11, r12], dim=-1),
        torch.stack([r20, r21, r22], dim=-1),
    ], dim=-2)
    
    if isinstance(quaternion, ndarray):
        return rotation_matrix.numpy()
    return rotation_matrix

def rotation_matrix_to_quaternion(rotation_matrix: Union[Tensor, ndarray]) -> Union[Tensor, ndarray]:
    r"""
    Convert rotation matrix to quaternion.

    Args:
        rotation_matrix (Union[Tensor, ndarray]): Rotation matrix with shape (..., 3, 3).

    Returns:
        Union[Tensor, ndarray], quaternion with shape (..., 4).
    """
    if isinstance(rotation_matrix, ndarray):
        rotation_matrix = torch.from_numpy(rotation_matrix)
    
    r00, r01, r02 = torch.unbind(rotation_matrix[..., 0, :], dim=-1)
    r10, r11, r12 = torch.unbind(rotation_matrix[..., 1, :], dim=-1)
    r20, r21, r22 = torch.unbind(rotation_matrix[..., 2, :], dim=-1)
    
    q0 = torch.sqrt(1 + r00 + r11 + r22) / 2
    q1 = (r21 - r12) / (4 * q0)
    q2 = (r02 - r20) / (4 * q0)
    q3 = (r10 - r01) / (4 * q0)
    
    quaternion = torch.stack([q0, q1, q2, q3], dim=-1)
    
    if isinstance(rotation_matrix, ndarray):
        return quaternion.numpy()
    return quaternion

def quaternion_multiply(q1: Union[Tensor, ndarray], q2: Union[Tensor, ndarray]) -> Union[Tensor, ndarray]:
    r"""
    Multiply two quaternions.

    Args:
        q1 (Union[Tensor, ndarray]): First quaternion with shape (..., 4).
        q2 (Union[Tensor, ndarray]): Second quaternion with shape (..., 4).

    Returns:
        Union[Tensor, ndarray], product quaternion with shape (..., 4).
    """
    if isinstance(q1, ndarray):
        q1 = torch.from_numpy(q1)
    if isinstance(q2, ndarray):
        q2 = torch.from_numpy(q2)
    
    w1, x1, y1, z1 = torch.unbind(q1, dim=-1)
    w2, x2, y2, z2 = torch.unbind(q2, dim=-1)
    
    w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
    x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
    y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
    z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
    
    product = torch.stack([w, x, y, z], dim=-1)
    
    if isinstance(q1, ndarray) and isinstance(q2, ndarray):
        return product.numpy()
    return product

def quaternion_conjugate(quaternion: Union[Tensor, ndarray]) -> Union[Tensor, ndarray]:
    r"""
    Get conjugate of quaternion.

    Args:
        quaternion (Union[Tensor, ndarray]): Quaternion with shape (..., 4).

    Returns:
        Union[Tensor, ndarray], conjugate quaternion with shape (..., 4).
    """
    if isinstance(quaternion, ndarray):
        quaternion = torch.from_numpy(quaternion)
    
    w, x, y, z = torch.unbind(quaternion, dim=-1)
    conjugate = torch.stack([w, -x, -y, -z], dim=-1)
    
    if isinstance(quaternion, ndarray):
        return conjugate.numpy()
    return conjugate

def quaternion_inverse(quaternion: Union[Tensor, ndarray]) -> Union[Tensor, ndarray]:
    r"""
    Get inverse of quaternion.

    Args:
        quaternion (Union[Tensor, ndarray]): Quaternion with shape (..., 4).

    Returns:
        Union[Tensor, ndarray], inverse quaternion with shape (..., 4).
    """
    if isinstance(quaternion, ndarray):
        quaternion = torch.from_numpy(quaternion)
    
    conjugate = quaternion_conjugate(quaternion)
    norm_squared = torch.sum(quaternion * quaternion, dim=-1, keepdim=True)
    inverse = conjugate / norm_squared
    
    if isinstance(quaternion, ndarray):
        return inverse.numpy()
    return inverse

def quaternion_normalize(quaternion: Union[Tensor, ndarray]) -> Union[Tensor, ndarray]:
    r"""
    Normalize quaternion.

    Args:
        quaternion (Union[Tensor, ndarray]): Quaternion with shape (..., 4).

    Returns:
        Union[Tensor, ndarray], normalized quaternion with shape (..., 4).
    """
    if isinstance(quaternion, ndarray):
        quaternion = torch.from_numpy(quaternion)
    
    norm = torch.norm(quaternion, dim=-1, keepdim=True)
    normalized = quaternion / norm
    
    if isinstance(quaternion, ndarray):
        return normalized.numpy()
    return normalized

def quaternion_rotate(point: Union[Tensor, ndarray], quaternion: Union[Tensor, ndarray]) -> Union[Tensor, ndarray]:
    r"""
    Rotate point by quaternion.

    Args:
        point (Union[Tensor, ndarray]): Point with shape (..., 3).
        quaternion (Union[Tensor, ndarray]): Quaternion with shape (..., 4).

    Returns:
        Union[Tensor, ndarray], rotated point with shape (..., 3).
    """
    if isinstance(point, ndarray):
        point = torch.from_numpy(point)
    if isinstance(quaternion, ndarray):
        quaternion = torch.from_numpy(quaternion)
    
    # Convert point to quaternion
    point_quat = torch.cat([torch.zeros_like(point[..., :1]), point], dim=-1)
    
    # Rotate point
    rotated = quaternion_multiply(
        quaternion_multiply(quaternion, point_quat),
        quaternion_conjugate(quaternion)
    )
    
    # Extract rotated point
    rotated_point = rotated[..., 1:]
    
    if isinstance(point, ndarray) and isinstance(quaternion, ndarray):
        return rotated_point.numpy()
    return rotated_point
